import os
import cv2
import numpy as np
from math import pi, acos
from joblib import Parallel, delayed, parallel_backend
from datetime import datetime


# -----     3.3 使用MIRD的多模态余弦相似性度量     ----- #

def cosine_similarity(a, b):
    """
    计算两个向量之间的余弦相似性。

    :param a: 向量a。
    :param b: 向量b。
    :return: 余弦相似性值。
    """
    dot_product = np.dot(a, b)
    norm_a = np.linalg.norm(a)
    norm_b = np.linalg.norm(b)

    if norm_a == 0 and norm_b == 0:
        return 0.0
    elif norm_a == 0 or norm_b == 0:
        return pi / 2
    else:
        cosine_value = dot_product / (norm_a * norm_b)
        cosine_value = np.clip(cosine_value, -1.0, 1.0)
        return acos(cosine_value)


def compute_similarity(y, x, ref_mird, temp_mird, ht, wt):
    """
    计算参考图像和模板图像在特定位置的相似性和累加。

    :param y: 行号。
    :param x: 列号。
    :param ref_mird: 参考图像的MIRD。
    :param temp_mird: 模板图像的MIRD。
    :param ht: 模板图像的高度。
    :param wt: 模板图像的宽度。
    :return: 相似性值和位置。
    """
    window = ref_mird[y:y + ht, x:x + wt]
    similarity_sum = 0.0
    for i in range(ht):
        for j in range(wt):
            similarity_sum += cosine_similarity(window[i, j], temp_mird[i, j])

    return similarity_sum, (y, x)


def multi_modal_cosine_similarity(ref_mird, temp_mird, step_size=1, n_jobs=-1):
    """
    使用余弦相似性度量计算多模态图像之间的相似性，并行处理。
    :param ref_mird: 参考图像的MIRD。
    :param temp_mird: 模板图像的MIRD。
    :param step_size: 滑动窗口的步长，默认值为1，实现逐像素滑动。
    :param n_jobs: 并行处理的作业数，默认值为-1，使用所有可用的CPU核。
    :return: 相似性地图，形状为 (h, w)。
    """
    ht, wt, dt = temp_mird.shape
    hr, wr, dr = ref_mird.shape

    if dt != dr:
        raise ValueError("特征向量的维度不一致")

    # 优化后的 similarity_map 形状
    similarity_map_height = (hr - ht) // step_size + 1
    similarity_map_width = (wr - wt) // step_size + 1
    similarity_map = np.zeros((similarity_map_height, similarity_map_width))

    # 处理进度变量
    total_steps = similarity_map_height * similarity_map_width
    processed_steps = 0

    def progress_callback(result):
        nonlocal processed_steps
        processed_steps += 1
        if processed_steps % 100 == 0:  # 每处理100步打印一次进度
            current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            print(f"处理进度：{processed_steps}/{total_steps} at {current_time}")
        return result

    def compute_and_update_similarity_map(y, x):
        similarity_sum, position = compute_similarity(y, x, ref_mird, temp_mird, ht, wt)
        similarity_map[y // step_size, x // step_size] = similarity_sum
        return progress_callback((similarity_sum, position))

    with parallel_backend('threading', n_jobs=n_jobs):
        results = Parallel()(
            delayed(compute_and_update_similarity_map)(y, x)
            for y in range(0, hr - ht + 1, step_size)
            for x in range(0, wr - wt + 1, step_size)
        )

    return similarity_map


def find_best_match_position(similarity_map):
    """
    找到最佳匹配位置。

    :param similarity_map: 相似性地图。
    :return: 最佳匹配位置。
    """
    best_match_position = np.unravel_index(np.argmin(similarity_map), similarity_map.shape)
    return best_match_position


def save_matching_image(ref_img, best_match_position, template_size, output_filename, step_size=1):
    """
    保存带有红色标记线的匹配图像。

    :param ref_img: 参考图像。
    :param best_match_position: 最佳匹配位置。
    :param template_size: 模板图像的大小。
    :param step_size: 滑动窗口的步长。
    :param output_filename: 输出图像文件名。
    """
    y, x = best_match_position
    y *= step_size
    x *= step_size
    h, w = template_size

    result_img = ref_img.copy()

    # 在最佳匹配位置画红色矩形，opencv里面不是RGB，而是BGR
    cv2.rectangle(result_img, (x, y), (x + w, y + h), (0, 0, 255), 2)

    # 创建输出目录
    output_path = os.path.join('../data/optical_to_SAR/matching', output_filename)
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    cv2.imwrite(output_path, result_img)
